import torch
from netualbuild import NeuralNetwork

from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
import visdom
import torch.nn as nn

device=(
  "cuda:0"
  if torch.cuda.is_available()
  else "mps"
  if torch.backends.mps.is_available()
  else "cpu"
)

print(f"device: {device}")

model = NeuralNetwork().to(device)
print(model)

labels_map = {
    0: "T-Shirt",
    1: "Trouser",    
    2: "Pullover",    
    3: "Dress",    
    4: "Coat",    
    5: "Sandal",    
    6: "Shirt",    
    7: "Sneaker",    
    8: "Bag",    
    9: "Ankle Boot",
}

training_data = datasets.FashionMNIST(root="data",train=True,download=True,transform=ToTensor())
test_data = datasets.FashionMNIST(root="data", train=False, download=True,transform=ToTensor())

train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

learning_rate = 1e-3
batch_size = 64
loss_fn = nn.CrossEntropyLoss()
optmizer = torch.optim.Adam(model.parameters(),lr=learning_rate)
# optmizer = torch.optim.SGD(model.parameters(),lr=learning_rate)

#构建tensorboard句柄
vis = visdom.Visdom(env='vis_fashion')

# props = {
#     'learning_rate': {'type': 'number', 'min': 0.0001, 'max': 0.1, 'step': 0.001},
#     'batch_size': {'type': 'number', 'min': 1, 'max': 64},
#     'optimizer': {'type': 'select', 'values': ['Adam', 'SGD']},
#     'reset_model': {'type': 'button', 'value': 'Reset'}
# }

# vis.properties(props, win='param_controls')  # 显示控件窗口

import gradio as gr

def update_model(learning_rate, batch_size):
    # 更新模型参数
    for param_group in optimizer.param_groups:
        param_group['lr'] = learning_rate
    return f"Updated: LR={learning_rate}, Batch={batch_size}"

iface = gr.Interface(
    fn=update_model,
    inputs=[gr.Slider(0.001, 0.1), gr.Number(32)],
    outputs="text"
)
iface.launch()

# for epochindex in range(0,2):
#   model.train()
#   for batchindex,(image,label) in enumerate(train_dataloader):
#     image = image.to(device)
#     label = label.to(device)
#     pred = model(image)
#     loss = loss_fn(pred,label)
#     # vis.text("Hello!")
#     vis.line(Y=[loss.item()],X=[batchindex],win='loss_plot',update='append',opts={'title':'TraingLoss'})
#     # vis.image(image.numpy(),win='input')
#     loss.backward()
#     optmizer.step()
#     optmizer.zero_grad()
#   model.eval()
#   batchnum = len(test_dataloader)
#   datasize = len(test_data)
#   test_loss = 0
#   test_correct = 0
#   with torch.no_grad():
#     for batchindex,(image,label) in enumerate(test_dataloader):
#       image = image.to(device)
#       label = label.to(device)
#       pred = model(image)
#       test_loss += loss_fn(pred,label)
#       # print(f"pred.shape:{pred.shape}  label.shape:{label.shape}\n")
#       test_correct += (pred.argmax(1) == label).type(torch.float).sum().item()
#       # print(f"pred:{pred.shape} label:{label.shape}, pred[0].size():{pred.size(0)}")
#       if torch.equal(pred.argmax(1), label):
#         for index in range(0,pred.size(0)):
#           predvalue = pred.argmax(1)
#           # print(predvalue,label)
#           print(f"pred:{labels_map[predvalue[index].item()]} label:{labels_map[label[index].item()]}")
#     test_loss /= batchnum
#     test_correct /= datasize
#     print(f"accuacy:{(100*test_correct):>0.1f}% avgloss:{test_loss:>8f}\n")